#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import tempfile
import re
import logging
import uuid
from typing import Dict, List, Any, BinaryIO, Tuple, Optional
from pathlib import Path
import shutil
import io
from datetime import datetime

from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.config.enums import SupportedPdfParseMethod

from document_processor.engines.base_engine import BaseEngine

logger = logging.getLogger(__name__)

class MinerUEngine(BaseEngine):
    """使用 MinerU 处理文档的引擎
    
    这是一个高级文档处理引擎，基于 OpenDataLab 的 MinerU 项目。
    支持 PDF、Word、PPT 等多种格式的文档转换，具有更好的布局分析和内容提取能力。
    """
    
    def __init__(self, config: Dict[str, Any]):
        """初始化 MinerU 处理引擎
        
        Args:
            config: 配置字典，包含 MinerU 相关的配置
        """
        super().__init__(config)
        
        # 初始化 MinerU 配置
        self.mineru_config = config.get("mineru", {})
        self.model_path = self.mineru_config.get("model_path", "models")
        logger.info("已初始化 MinerU 处理引擎")

    async def convert_to_markdown(
        self, 
        file_content: BinaryIO, 
        file_name: str,
        image_option: str = "obs",
        work_dir: Optional[str] = None,
        md_filename: Optional[str] = None
    ) -> Dict[str, Any]:
        """使用 MinerU 将文档转换为 Markdown 格式
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            image_option: 图片处理选项，可选值："obs"（上传到OBS）或"remove"（移除图片）
            work_dir: 工作目录路径，如果提供则使用此目录，否则创建临时目录
            md_filename: 生成的Markdown文件名（不含扩展名），用于命名图片目录
            
        Returns:
            包含以下键的字典:
                - text: 提取的纯文本
                - markdown: 转换后的Markdown文本
                - images: 图片信息列表
                - temp_dir: 临时目录路径
        """
        # 验证图片处理选项
        if image_option not in ["obs", "remove"]:
            logger.warning(f"不支持的图片处理选项: {image_option}，使用 obs 选项")
            image_option = "obs"
        
        # 检查文件格式
        file_ext = os.path.splitext(file_name)[1].lower()
        supported_formats = ['.pdf']  # 目前仅支持PDF
        if file_ext not in supported_formats:
            raise ValueError(f"MinerU引擎目前仅支持PDF格式，不支持: {file_ext}")
            
        # 使用提供的工作目录或创建临时目录
        temp_dir = work_dir if work_dir else tempfile.mkdtemp()
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
            
        # 确保output目录存在
        output_dir = os.path.join(temp_dir, "output")
        local_image_dir = os.path.join(output_dir, "images")
        os.makedirs(local_image_dir, exist_ok=True)
            
        temp_file_path = os.path.join(temp_dir, file_name)
        name_without_suff = Path(file_name).stem
        
        try:
            # 保存上传的文件到临时位置
            file_content_bytes = file_content.read()
            with open(temp_file_path, "wb") as f:
                f.write(file_content_bytes)
            
            # 创建图片目录名称
            if md_filename:
                images_dir_name = f"{md_filename}_images"
            else:
                dir_id = uuid.uuid4().hex[:8]
                images_dir_name = f"doc_images_{dir_id}"
                
            # 设置数据写入器
            image_writer = FileBasedDataWriter(local_image_dir)
            md_writer = FileBasedDataWriter(output_dir)
            
            logger.info(f"开始使用 MinerU 转换文档: {file_name}")
            
            # 创建数据集实例
            ds = PymuDocDataset(file_content_bytes)
            
            # 根据文档类型选择处理方式
            if ds.classify() == SupportedPdfParseMethod.OCR:
                infer_result = ds.apply(doc_analyze, ocr=True)
                pipe_result = infer_result.pipe_ocr_mode(image_writer)
            else:
                infer_result = ds.apply(doc_analyze, ocr=False)
                pipe_result = infer_result.pipe_txt_mode(image_writer)
            
            # 生成分析结果文件
            infer_result.draw_model(os.path.join(output_dir, f"{name_without_suff}_model.pdf"))
            pipe_result.draw_layout(os.path.join(output_dir, f"{name_without_suff}_layout.pdf"))
            pipe_result.draw_span(os.path.join(output_dir, f"{name_without_suff}_spans.pdf"))
            
            # 获取Markdown内容
            markdown_text = pipe_result.get_markdown(images_dir_name)
            
            # 保存Markdown和其他结果
            pipe_result.dump_md(md_writer, f"{name_without_suff}.md", images_dir_name)
            pipe_result.dump_content_list(md_writer, f"{name_without_suff}_content_list.json", images_dir_name)
            pipe_result.dump_middle_json(md_writer, f"{name_without_suff}_middle.json")
            
            # 收集图片信息
            images = []
            image_files = [f for f in os.listdir(local_image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
            
            for idx, img_file in enumerate(image_files):
                try:
                    img_path = os.path.join(local_image_dir, img_file)
                    with open(img_path, "rb") as f:
                        image_data = f.read()
                    
                    img_format = Path(img_file).suffix.lstrip(".")
                    if not img_format:
                        img_format = "png"
                    
                    image_info = {
                        "filename": img_file,
                        "format": img_format,
                        "data": image_data,
                        "dir_name": images_dir_name
                    }
                    images.append(image_info)
                    
                except Exception as e:
                    logger.error(f"处理图片失败: {str(e)}")
                    continue
            
            return {
                "text": re.sub(r'!\[.*?\]\(.*?\)', '', markdown_text),  # 纯文本（移除图片引用）
                "markdown": markdown_text,
                "images": images,
                "temp_dir": temp_dir,
                "output_md": os.path.join(output_dir, f"{name_without_suff}.md"),
                "images_dir": local_image_dir,
                "md_filename": name_without_suff
            }
            
        except Exception as e:
            logger.error(f"处理文件出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            raise

    async def handle_images(
        self, 
        images: List[Dict[str, Any]], 
        markdown_text: str,
        uploader,
        md_filename: Optional[str] = None
    ) -> str:
        """处理图片（上传到OBS并更新Markdown中的图片引用）
        
        Args:
            images: 图片信息列表
            markdown_text: Markdown文本
            uploader: OBS上传器实例
            md_filename: Markdown文件名（不含扩展名），用于生成OBS路径
            
        Returns:
            更新后的Markdown文本（图片引用已替换为OBS URL）
        """
        if not images:
            logger.info("没有图片需要处理")
            return markdown_text
        
        if not uploader or not hasattr(uploader, 'upload'):
            logger.warning("未提供有效的OBS上传器，无法处理图片")
            return markdown_text
        
        logger.info(f"开始处理 {len(images)} 张图片")
        
        try:
            # 解析所有占位符
            placeholder_pattern = r'__IMAGE_URL_(\d+)__'
            all_placeholders = re.findall(placeholder_pattern, markdown_text)
            logger.info(f"在Markdown中找到 {len(all_placeholders)} 个图片占位符")
            
            # 为每张图片获取URL（上传到OBS）
            image_urls = {}
            
            # 生成OBS路径前缀
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            obs_path_prefix = ""
            if md_filename:
                safe_filename = re.sub(r'[^\w\-\.]', '_', md_filename)
                obs_path_prefix = f"docs/{timestamp}/{safe_filename}"
            else:
                obs_path_prefix = f"docs/{timestamp}/document"
            logger.info(f"使用OBS路径前缀: {obs_path_prefix}")
            
            for i, image in enumerate(images):
                try:
                    if not image or "data" not in image:
                        logger.warning(f"跳过无效的图片: {image.get('filename', '未知')}")
                        continue
                        
                    logger.info(f"处理图片 {i+1}/{len(images)}: {image['filename']}")
                    
                    # 准备上传数据
                    upload_data = {
                        "filename": image['filename'],
                        "data": image['data'],
                        "content_type": f"image/{image['format']}"
                    }
                    
                    # 构建OBS路径
                    safe_doc_name = re.sub(r'[^\w\-\.]', '_', md_filename if md_filename else "doc")
                    new_filename = f"{safe_doc_name}_img{i+1:03d}.{image['format']}"
                    obs_path = f"{obs_path_prefix}/{new_filename}"
                    upload_data["object_key"] = obs_path
                    
                    # 上传图片
                    img_url = await uploader.upload(upload_data)
                    if not img_url:
                        logger.warning(f"图片 {image['filename']} 上传返回空URL")
                        continue
                        
                    # 记录URL
                    placeholder = f"__IMAGE_URL_{i}__"
                    image_urls[placeholder] = img_url
                    logger.info(f"图片上传成功: {img_url}")
                    
                except Exception as e:
                    logger.error(f"处理图片 {image.get('filename', '未知')} 时出错: {str(e)}")
                    continue
            
            # 替换所有占位符
            def replace_placeholder(match):
                placeholder = match.group(0)
                if placeholder not in image_urls:
                    return placeholder
                
                url = image_urls[placeholder]
                img_idx = int(match.group(1))
                img_caption = f"图片_{img_idx+1}"
                
                return f"![{img_caption}]({url})"
            
            processed_markdown = re.sub(placeholder_pattern, replace_placeholder, markdown_text)
            
            # 清理未替换的占位符
            lines = processed_markdown.split('\n')
            filtered_lines = []
            
            for line in lines:
                if re.match(r'^\s*__IMAGE_URL_\d+__\s*$', line):
                    continue
                filtered_lines.append(line)
            
            processed_markdown = '\n'.join(filtered_lines)
            
            logger.info(f"图片处理完成，共处理 {len(images)} 张图片")
            return processed_markdown
            
        except Exception as e:
            logger.error(f"处理图片集合时出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            return markdown_text 

    async def extract_images(
        self,
        file_content: BinaryIO,
        file_name: str,
        work_dir: Optional[str] = None
    ) -> Dict[str, Any]:
        """从文档中提取图片
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            work_dir: 工作目录路径，如果提供则使用此目录，否则创建临时目录
            
        Returns:
            包含以下键的字典:
                - images: 图片信息列表
                - temp_dir: 临时目录路径
        """
        # 使用提供的工作目录或创建临时目录
        temp_dir = work_dir if work_dir else tempfile.mkdtemp()
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
            
        # 确保output目录存在
        output_dir = os.path.join(temp_dir, "output")
        local_image_dir = os.path.join(output_dir, "images")
        os.makedirs(local_image_dir, exist_ok=True)
            
        temp_file_path = os.path.join(temp_dir, file_name)
        name_without_suff = Path(file_name).stem
        
        try:
            # 保存上传的文件到临时位置
            file_content_bytes = file_content.read()
            with open(temp_file_path, "wb") as f:
                f.write(file_content_bytes)
            
            # 设置数据写入器
            image_writer = FileBasedDataWriter(local_image_dir)
            
            # 创建数据集实例
            ds = PymuDocDataset(file_content_bytes)
            
            # 根据文档类型选择处理方式并提取图片
            if ds.classify() == SupportedPdfParseMethod.OCR:
                infer_result = ds.apply(doc_analyze, ocr=True)
                pipe_result = infer_result.pipe_ocr_mode(image_writer)
            else:
                infer_result = ds.apply(doc_analyze, ocr=False)
                pipe_result = infer_result.pipe_txt_mode(image_writer)
            
            # 收集图片信息
            images = []
            image_files = [f for f in os.listdir(local_image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
            
            for idx, img_file in enumerate(image_files):
                try:
                    img_path = os.path.join(local_image_dir, img_file)
                    with open(img_path, "rb") as f:
                        image_data = f.read()
                    
                    img_format = Path(img_file).suffix.lstrip(".")
                    if not img_format:
                        img_format = "png"
                    
                    image_info = {
                        "filename": img_file,
                        "format": img_format,
                        "data": image_data,
                        "path": img_path
                    }
                    images.append(image_info)
                    
                except Exception as e:
                    logger.error(f"处理图片失败: {str(e)}")
                    continue
            
            return {
                "images": images,
                "temp_dir": temp_dir
            }
            
        except Exception as e:
            logger.error(f"提取图片时出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            raise 